# -*- coding: utf-8 -*-
"""
@Time: 2021/11/23 17:20
@Author: 鹄望潇湘
@File: train.py
@desc: 
"""
from torch.nn import Module
from torch.utils.data.dataset import Dataset
from tools_zsj.Trainer import Trainer, TrainParameter
import torch
from dataset.GaitMChannelDataSet import GaitMChannelDataLoader, GaitDataSet
from backbone.vgg_16 import VGG16
from backbone.MyGaitNet import CoAtGait
from loss.metrics import MagFace, ArcFace
from loss.focal import FocalLoss


class GaitTrainer(Trainer):

    def __init__(self, module: Module, dataset: Dataset, parameter: TrainParameter):
        super(GaitTrainer, self).__init__(module, dataset, parameter)
        self.optimizer = torch.optim.SGD(self.module.parameters(), lr=parameter.get_parameter("lr"),
                                         momentum=0.009, dampening=0.0001)
        self.head = ArcFace(2048, 124, [0])
        self.loss_function = FocalLoss()
        self.module.to(parameter.get_parameter('device'))
        self.sum_batch = len(dataset)//parameter.batch_size


    def train_detail(self, index: int, item: list):
        self.optimizer.zero_grad()
        id = item[1]
        device = self.train_parameter.get_parameter("device")
        x = item[0].float().to(device)
        output1 = self.module(x)
        output2 = self.head(output1, id.to(device))
        loss = self.loss_function(output2, id.to(device))
        loss.backward()
        self.optimizer.step()
        if index % 20 == 0:
            self.logger.info("batch {}/{} train loss: {:.3f}".format(index, self.sum_batch, loss.item()))

        pass

    def epoch_finish(self, epoch: int):
        self.logger.info("第{}个epoch已经训练完成".format(epoch))
        if (epoch+1) % 20 == 0:
            torch.save(self.module, "checkpoint_"+str(epoch)+".pth")
        pass

    def train_finish(self):
        self.logger.info("训练已完成")
        pass


if __name__=='__main__':
    data_loader = GaitMChannelDataLoader("E:\\DataSet\\CASIA-B-preprocessed", is_preprocessed=True, use_cache=True)
    # datas = data_loader.get_dataset_with_type_and_angle("nm-02", "090")
    datas = data_loader.get_partial_data_with_id([index for index in range(0, 75)])
    dataset = GaitDataSet(datas)
    net = CoAtGait(30, 64)
    parameter = TrainParameter(1000, 64, 0, True)
    parameter.add_parameter("lr", 0.01)
    parameter.add_parameter("device", torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'))
    trainer = GaitTrainer(net, dataset, parameter)
    trainer.start_train()
